//
// Created by whitby on 8/4/22.
//

#include <fstream>
#include "surfel_map.h"

using namespace std;

SurfelMap::SurfelMap(int width, int height, Eigen::Matrix3f K, float far_dist, float near_dist, int drift_free_poses) {
    cam_width = width;
    cam_height = height;
    camera_matrix = std::move(K);
    cam_fx = camera_matrix(0, 0);
    cam_cx = camera_matrix(0, 2);
    cam_fy = camera_matrix(1, 1);
    cam_cy = camera_matrix(1, 2);
    this->drift_free_poses = drift_free_poses;

    printf("Have the following settings: \n");
    printf("camera matrix: \n");
    cout << camera_matrix << endl;
    printf("fuse the distance between %4f m and %4f m. \n", near_dist, far_dist);
    fusion_functions.initialize(cam_width, cam_height, cam_fx, cam_fy, cam_cx, cam_cy,
                                near_dist, far_dist);
}

SurfelMap::~SurfelMap() = default;

void SurfelMap::orb_result_process(Eigen::Matrix4f &received_pose, int reference_index,
                                   bool is_new_keyFrame, vector<Eigen::Matrix4f> &allKeyFrames,
                                   vector<pair<int, int>> &loops) {
    static Eigen::Matrix4f transform_kitti;
    {
        if (poses_database.empty()) {
            Eigen::Matrix4f idea_pose = Eigen::Matrix4f::Identity();
            transform_kitti = idea_pose * received_pose.inverse();
        }
    }
    Eigen::Matrix4f transform_pose = transform_kitti * received_pose;

    // first update to the poses
    bool loop_changed = false;
    for (size_t i = 0; i < poses_database.size() && i < allKeyFrames.size(); ++i) {
        Eigen::Matrix4f after_pose = transform_kitti * allKeyFrames[i];
        poses_database[i].loop_pose = &after_pose;
        if (poses_database[i].loopChanged()) {
            loop_changed = true;
        }
    }

    if (poses_database.size() > allKeyFrames.size()) {
        size_t last_update_index = allKeyFrames.size() - 1;
        size_t start_index = allKeyFrames.size();


        Eigen::Matrix4f warp_pose = poses_database[last_update_index].getWarpPose();

        for (; start_index < poses_database.size(); ++start_index) {
            Eigen::Matrix4f after_pose = warp_pose * (*poses_database[start_index].cam_pose);
            poses_database[start_index].loop_pose = &after_pose;
        }
    }

    // warp the surfels
    if (loop_changed)
        warp_surfels();

    // add loop information
    size_t loop_num = loops.size();
    for (size_t i = 0; i < loop_num; ++i) {
        int loop_first = loops[i].first;
        int loop_second = loops[i].second;
        if (loop_first < poses_database.size() && loop_second < poses_database.size()) {
            if (std::find(
                    poses_database[loop_first].linked_pose_index.begin(),
                    poses_database[loop_first].linked_pose_index.end(),
                    loop_second) == poses_database[loop_first].linked_pose_index.end()) {
                poses_database[loop_first].linked_pose_index.push_back(loop_second);
            }
            if (std::find(
                    poses_database[loop_second].linked_pose_index.begin(),
                    poses_database[loop_second].linked_pose_index.end(),
                    loop_first) == poses_database[loop_second].linked_pose_index.end()) {
                poses_database[loop_second].linked_pose_index.push_back(loop_first);
            }
        } else {
            printf("cannot find loop pose %d and %d, we have %d poses!\n",
                   loop_first, loop_second, (int) poses_database.size());
        }
    }

    if (poses_database.empty()) {
        is_new_keyFrame = true; // the first is a keyframe, but ORB_SLAM2.getKeyFrame return false
    }

    if (is_new_keyFrame) {
        // add new poses
        printf("add new poses \n");
        PoseElement pose_element;
        int this_pose_index = static_cast<int>(poses_database.size());
        pose_element.cam_pose = &transform_pose;
        pose_element.loop_pose = &transform_pose;
        if (!poses_database.empty()) {
            pose_element.linked_pose_index.push_back(reference_index);
            poses_database[reference_index].linked_pose_index.push_back(this_pose_index);
        }
        poses_database.push_back(pose_element);
        local_surfels_indexs.insert(this_pose_index);
    }
    cout << "the size of poses database is " << static_cast<int>(poses_database.size()) << endl;
}

void SurfelMap::warp_surfels() {
    warp_thread_pool.clear();
    warp_thread_num = 10;

    // warp inactive surfels
    for (int i = 0; i < warp_thread_num; ++i) {
        thread this_thread(&SurfelMap::warp_inactive_surfels_cpu_kernel, this, i, warp_thread_num);
        warp_thread_pool.push_back(std::move(this_thread));
    }

    // warp active surfels
    int local_index = *local_surfels_indexs.begin();
    Eigen::Matrix4f warp_pose = poses_database[local_index].getWarpPose();

    for (int i = 0; i < warp_thread_num; ++i) {
        thread this_thread(&SurfelMap::warp_activate_surfels_cpu_kernel, this, i, warp_thread_num, warp_pose);
        warp_thread_pool.push_back(std::move(this_thread));
    }

    for (auto &thread: warp_thread_pool) {
        if (thread.joinable()) {
            thread.join();
        }
    }
}

void SurfelMap::warp_inactive_surfels_cpu_kernel(int thread_i, int thread_num) {
    int step = static_cast<int>(poses_database.size()) / thread_num;
    int begin_index = step * thread_i;
    int end_index = begin_index + step;
    if (thread_i == thread_num - 1)
        end_index = static_cast<int>(poses_database.size());

    for (int i = begin_index; i < end_index; ++i) {
        if (!poses_database[i].loopChanged())
            continue;
        if (poses_database[i].attached_surfels.empty()) {
            poses_database[i].cam_pose = poses_database[i].loop_pose;
            continue;
        }

        Eigen::Matrix4f warp_matrix = poses_database[i].getWarpPose();
        Eigen::MatrixXf point_positions(4, poses_database[i].attached_surfels.size());
        Eigen::MatrixXf point_norms(3, poses_database[i].attached_surfels.size());
        for (int surfel_i = 0; surfel_i < poses_database[i].attached_surfels.size(); surfel_i++) {
            point_positions(0, surfel_i) = poses_database[i].attached_surfels[surfel_i]->px;
            point_positions(1, surfel_i) = poses_database[i].attached_surfels[surfel_i]->py;
            point_positions(2, surfel_i) = poses_database[i].attached_surfels[surfel_i]->pz;
            point_positions(3, surfel_i) = 1.0;
            point_norms(0, surfel_i) = poses_database[i].attached_surfels[surfel_i]->nx;
            point_norms(1, surfel_i) = poses_database[i].attached_surfels[surfel_i]->ny;
            point_norms(2, surfel_i) = poses_database[i].attached_surfels[surfel_i]->nz;
        }

        point_positions = warp_matrix * point_positions;
        point_norms = warp_matrix.block<3, 3>(0, 0) * point_norms;

        for (int surfel_i = 0; surfel_i < poses_database[i].attached_surfels.size(); surfel_i++) {
            poses_database[i].attached_surfels[surfel_i]->px = point_positions(0, surfel_i);
            poses_database[i].attached_surfels[surfel_i]->py = point_positions(1, surfel_i);
            poses_database[i].attached_surfels[surfel_i]->pz = point_positions(2, surfel_i);
            poses_database[i].attached_surfels[surfel_i]->nx = point_norms(0, surfel_i);
            poses_database[i].attached_surfels[surfel_i]->ny = point_norms(1, surfel_i);
            poses_database[i].attached_surfels[surfel_i]->nz = point_norms(2, surfel_i);
        }

        poses_database[i].cam_pose = poses_database[i].loop_pose;
    }
}

void SurfelMap::warp_activate_surfels_cpu_kernel(int thread_i, int thread_num, const Eigen::Matrix4f &transform_m) {
    int step = static_cast<int>(local_surfels.size()) / thread_num;
    int begin_index = step * thread_i;
    int end_index = begin_index + step;
    if (thread_i == thread_num - 1)
        end_index = static_cast<int>(local_surfels.size());
    int surfel_num = end_index - begin_index;

    Eigen::MatrixXf point_positions(4, surfel_num);
    Eigen::MatrixXf point_norms(3, surfel_num);

    for (int i = 0; i < surfel_num; ++i) {
        point_positions(0, i) = local_surfels[i + begin_index]->px;
        point_positions(1, i) = local_surfels[i + begin_index]->py;
        point_positions(2, i) = local_surfels[i + begin_index]->pz;
        point_positions(3, i) = 1.0;
        point_norms(0, i) = local_surfels[i + begin_index]->nx;
        point_norms(1, i) = local_surfels[i + begin_index]->ny;
        point_norms(2, i) = local_surfels[i + begin_index]->nz;
    }

    point_positions = transform_m * point_positions;
    point_norms = transform_m.block<3, 3>(0, 0) * point_norms;

    for (int i = 0; i < surfel_num; ++i) {
        local_surfels[i + begin_index]->px = point_positions(0, i);
        local_surfels[i + begin_index]->py = point_positions(1, i);
        local_surfels[i + begin_index]->pz = point_positions(2, i);
        local_surfels[i + begin_index]->nx = point_norms(0, i);
        local_surfels[i + begin_index]->ny = point_norms(1, i);
        local_surfels[i + begin_index]->nz = point_norms(2, i);
    }
}

void SurfelMap::synchronize(cv::Mat &left_img, cv::Mat &depth_img, int reference_index, Eigen::Matrix4f &current_pose) {
    // extract local surfels
    move_add_surfels(reference_index);

    cout << "the size of local_surfels is " << local_surfels.size() << endl;

    // fuse map
    fuse_map(left_img, depth_img, current_pose, reference_index);
}

void SurfelMap::move_add_surfels(int reference_index) {
    vector<int> poses_to_add, poses_to_remove;
    get_add_remove_poses(reference_index, poses_to_add, poses_to_remove);

    // remove surfels from "local_surfels" to "pose_database[i].attach_surfels"
    int removes = 0;
    if (!poses_to_remove.empty()) {
        for (int inactive_index: poses_to_remove) {
            for (auto &surfel: local_surfels) {
                if (surfel->update_times > 0 && surfel->last_update == inactive_index) {
                    poses_database[inactive_index].attached_surfels.push_back(surfel);
                    surfel->update_times = 0;
                    removes += 0;
                }
            }
            local_surfels_indexs.erase(inactive_index);
        }
    }
    cout << "remove " << removes << " surfels from local surfels to poses database." << endl;

    size_t adds = 0;
    if (!poses_to_add.empty()) {
        // add indexs
        local_surfels_indexs.insert(poses_to_add.begin(), poses_to_add.end());
        // add surfels to local
        for (int add_index: poses_to_add) {
            local_surfels.insert(
                    local_surfels.end(),
                    poses_database[add_index].attached_surfels.begin(),
                    poses_database[add_index].attached_surfels.end());
            adds += poses_database[add_index].attached_surfels.size();
            poses_database[add_index].attached_surfels.clear();
        }
    }
    cout << "add " << adds << " surfels from poses database to local surfels." << endl;
}

void SurfelMap::get_add_remove_poses(int root_index, vector<int> &pose_to_add, vector<int> &pose_to_remove) {
    vector<int> driftFree_poses;
    get_drift_free_poses(root_index, driftFree_poses, drift_free_poses);

    cout << "drift free poses: ";
    for (int pose_index : driftFree_poses) {
        cout << pose_index << " ";
    }
    cout << endl;

    cout << "local surfels index: ";
    for (int index : local_surfels_indexs) {
        cout << index << " ";
    }
    cout << endl;

    pose_to_add.clear();
    pose_to_remove.clear();

    // get to add
    for (int temp_pose_index: driftFree_poses) {
        if (local_surfels_indexs.find(temp_pose_index) == local_surfels_indexs.end())
            pose_to_add.push_back(temp_pose_index);
    }

    cout << "add " << pose_to_add.size() << " poses" << endl;

    // get to remove
    for (int surfels_pose_index: local_surfels_indexs) {
        if (std::find(driftFree_poses.begin(), driftFree_poses.end(), surfels_pose_index) == driftFree_poses.end()) {
            pose_to_remove.push_back(surfels_pose_index);
        }
    }

    cout << "remove " << pose_to_remove.size() << " poses" << endl;
}

void SurfelMap::get_drift_free_poses(int root_index, vector<int> &driftFree_poses, int drift_free_range) {
    if (poses_database.size() < root_index + 1) {
        printf("get_drift_free_poses: pose database do not have the root index!"
               "This should only happen in initialization!\n");
        return;
    }

    vector<int> this_level;
    vector<int> next_level;
    this_level.push_back(root_index);
    driftFree_poses.push_back(root_index);

    // get the drift
    for (int i = 1; i < drift_free_range; ++i) {
        for (int this_level_pose_index: this_level) {
            for (int linked_pose_index: poses_database[this_level_pose_index].linked_pose_index) {
                bool already_saved = (find(
                        driftFree_poses.begin(), driftFree_poses.end(), linked_pose_index
                ) != driftFree_poses.end());
                if (!already_saved) {
                    next_level.push_back(linked_pose_index);
                    driftFree_poses.push_back(linked_pose_index);
                }
            }
        }
        this_level.swap(next_level);
        next_level.clear();
    }
}

void SurfelMap::fuse_map(cv::Mat &image, cv::Mat &depth, Eigen::Matrix4f &fused_pose, int reference_index) {
    vector<shared_ptr<SurfelElement>> new_surfels;
    fusion_functions.fuse_initialize_map(
            reference_index, image, depth, fused_pose, local_surfels, new_surfels
    );

    // get the deleted surfel index
    vector<int> deleted_index;
    for (int i = 0; i < local_surfels.size(); ++i) {
        if (local_surfels[i]->update_times == 0)
            deleted_index.push_back(i);
    }

    cout << "delete surfels: " << (int) deleted_index.size() << endl;
    cout << "add new initialized surfels: " << (int) new_surfels.size() << endl;

    // add new initialized surfels
    for (auto &surfel: new_surfels) {
        if (surfel->update_times != 0) {
            if (!deleted_index.empty()) {
                int delete_surfel_index = deleted_index.back();
                int pose_index = local_surfels[delete_surfel_index]->last_update;
                poses_database[pose_index].attached_surfels.push_back(local_surfels[delete_surfel_index]);
                local_surfels[delete_surfel_index] = surfel;
            } else {
                local_surfels.push_back(surfel);
            }
        }
    }

    while (!deleted_index.empty()) {
        int delete_surfel_index = deleted_index.back();
        int pose_index = local_surfels[delete_surfel_index]->last_update;
        poses_database[pose_index].attached_surfels.push_back(local_surfels[delete_surfel_index]);
        local_surfels[delete_surfel_index] = local_surfels.back();
        deleted_index.pop_back();
        local_surfels.pop_back();
    }
}

void SurfelMap::save_cloud(const std::string &save_path) {
    printf("saving point cloud ...\n");
    int point_clouds = 0;
    for (auto &pose: poses_database) {
        point_clouds += static_cast<int>(pose.attached_surfels.size());
    }
    point_clouds += static_cast<int>(local_surfels.size());
    point_clouds += static_cast<int>(inactivate_surfels.size());

    cout << "the number of all point cloud is " << point_clouds << endl;

    std::ofstream stream(save_path.c_str());

    stream << "# .PCD v0.7 - Point Cloud Data file format" << std::endl;
    stream << "VERSION 0.7" << std::endl;
    stream << "FIELDS x y z intensity" << std::endl;
    stream << "SIZE 4 4 4 4" << std::endl;
    stream << "TYPE F F F F" << std::endl;
    stream << "COUNT 1 1 1 1" << std::endl;
    stream << "WIDTH " << point_clouds << std::endl;
    stream << "HEIGHT 1" << std::endl;
    stream << "VIEWPOINT 0 0 0 1 0 0 0" << std::endl;
    stream << "POINTS " << point_clouds << std::endl;
    stream << "DATA ascii" << std::endl;

    for (auto &surfel: local_surfels) {
        stream << surfel->px << " " << surfel->py << " " << surfel->pz << " " << surfel->color << endl;
    }

    for (auto &pose: poses_database) {
        for (auto &surfel: pose.attached_surfels) {
            stream << surfel->px << " " << surfel->py << " " << surfel->pz << " " << surfel->color << endl;
        }
    }

    for (auto &surfel : inactivate_surfels) {
        stream << surfel->px << " " << surfel->py << " " << surfel->pz << " " << surfel->color << endl;
    }

    stream.close();
    cout << "save point cloud done!" << endl;
}